/*
 *   Copyright 2023 <a href="mailto:asialjim@hotmail.com">Asial Jim</a>
 *
 *   Licensed under the Apache License, Version 2.0 (the "License");
 *   you may not use this file except in compliance with the License.
 *   You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *   See the License for the specific language governing permissions and
 *   limitations under the License.
 */
package com.asialjim.encommunicate.rsa.filter;

import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.IOUtils;

import java.io.ByteArrayOutputStream;
import java.io.Closeable;
import java.io.IOException;
import java.util.Objects;

/**
 * 基于  RSA 的Http响应包装
 *
 * @author Copyright © <a href="mailto:asialjim@hotmail.com">Asial Jim</a>   Co., LTD
 * @version 1.0
 * @since 2023/2/27, &nbsp;&nbsp; <em>version:1.0</em>, &nbsp;&nbsp; <em>java version:17</em>
 */
@Slf4j
public class SimpleRSAHttpServletResponseWrapper extends HttpServletResponseWrapper implements Closeable {

    private final ByteArrayOutputStream buffer;
    private final ServletOutputStream outputStream;

    public SimpleRSAHttpServletResponseWrapper(HttpServletResponse response) {
        super(response);

        this.buffer = new ByteArrayOutputStream();
        this.outputStream = new WrapperOutputStream(buffer);
    }

    @Override
    public ServletOutputStream getOutputStream() {
        return this.outputStream;
    }

    @Override
    public void flushBuffer() throws IOException {
        if (null != outputStream) {
            outputStream.flush();
        }
    }

    public byte[] getContent() throws IOException {
        flushBuffer();
        return buffer.toByteArray();
    }

    @Override
    public void close() {
        IOUtils.closeQuietly(this.buffer);
        IOUtils.closeQuietly(this.outputStream);
    }

    static class WrapperOutputStream extends ServletOutputStream {
        private final ByteArrayOutputStream bos;

        public WrapperOutputStream(ByteArrayOutputStream bos) {
            this.bos = bos;
        }

        @Override
        public void write(int b) {
            bos.write(b);
        }

        @Override
        public void write(@SuppressWarnings("NullableProblems") byte[] b) throws IOException {
            bos.write(b);
        }

        @Override
        public void write(@SuppressWarnings("NullableProblems") byte[] b, int off, int len) {
            bos.write(b, off, len);
        }

        @Override
        public boolean isReady() {
            return Objects.nonNull(bos);
        }

        @Override
        public void setWriteListener(WriteListener arg0) {
        }
    }
}